from langchain.agents import initialize_agent, AgentType
from langchain.tools import Tool

from a0base.base_llm import pop_llm

# 子 Agent 1：数学专家
math_agent = initialize_agent(
    tools=[],
    llm=pop_llm,
    agent=AgentType.OPENAI_FUNCTIONS,
    description="负责所有数学/代码问题"
)


# 子 Agent 2：搜索专家（假装有搜索工具）
def fake_search(q): return f"搜索《{q}》的结果：42"
search_tool = Tool(name="FakeSearch", func=fake_search, description="返回固定答案")
search_agent = initialize_agent(
    tools=[search_tool],
    llm=pop_llm,
    agent=AgentType.OPENAI_FUNCTIONS,
    description="负责实时信息检索"
)

# 用 Tool 包装子 Agent
math_tool = Tool(
    name="MathAgent",
    func=lambda q: math_agent.run(q),
    description="数学/逻辑/代码计算专用 Agent"
)
search_tool_wrapped = Tool(
    name="SearchAgent",
    func=lambda q: search_agent.run(q),
    description="实时信息检索 Agent"
)

# Supervisor Agent
supervisor = initialize_agent(
    tools=[math_tool, search_tool_wrapped],
    llm=pop_llm,
    agent=AgentType.OPENAI_FUNCTIONS
)

if __name__ == "__main__":
    question = "355/113 的近似精度是多少？再帮我搜一下『圆周率世界纪录』"
    print(supervisor.run(question))
